#include <mpi.h>
#include <iostream>
#include <vector>
#include <string>
#include <stdexcept>
#include <sstream>

class MPI_Exception : public std::runtime_error
{
public:
	MPI_Exception(std::string const& msg)
	 : std::runtime_error(msg)
	{ }
};

//#define ENFORCE(x) (x);
#define ENFORCE(x) { enforce(x, #x, __LINE__, __FILE__); }
void enforce(int error, std::string const& msg, long line, const char * file)
{
	if ( error != 0 )
	{
		std::stringstream ss;
		ss << "MPI Error: " << error << "\n";
		ss << "File: \"" << file << " - Line: " << line << "\n";
		ss << "Code: \"" << msg << "\"";
		throw MPI_Exception(ss.str());
	}
}

class MPI_Encapsulation
{
	std::size_t m_count;
	std::size_t m_number;
	
public:
	MPI_Encapsulation(int & argc, char ** argv[])
	{
		ENFORCE(MPI_Init(&argc, argv));
		int tmp;
		ENFORCE(MPI_Comm_size(MPI_COMM_WORLD, &tmp));
		this->m_count = static_cast<std::size_t>(tmp);
		ENFORCE(MPI_Comm_rank(MPI_COMM_WORLD, &tmp));
		this->m_number = static_cast<std::size_t>(tmp);
	}
	
	~MPI_Encapsulation()
	{
		int error = MPI_Finalize();
		if ( error )
			std::cerr << "Error: " << error << std::endl;
	}
	
	std::size_t getProcessCount() const
	{
		return this->m_count;
	}
	
	std::size_t getProcessNumber() const
	{
		return this->m_number;
	}
	
	bool isParentProcess() const
	{
		return this->getProcessNumber() == 0;
	}
	
public:
	template<class T>
	void receive(std::vector<T> & buffer, std::size_t sender) const
	{
		MPI_Status status;
		ENFORCE(MPI_Recv(&buffer[0], buffer.size(), MPI_INT, sender, 0, MPI_COMM_WORLD, &status)); /*todo: "MPI_INT" eventuell aendern, je nach template parameter*/
	}
	
	template<class T>
	void send(std::vector<T> const& buffer, std::size_t receiver) const
	{
		ENFORCE(MPI_Send(const_cast<T*>(&buffer[0]), buffer.size(), MPI_INT, receiver, 0, MPI_COMM_WORLD));
	}
	
	/*void scatter(std::size_t count, std::size_t target)
	{
		MPI_Scatter(nullptr, count, MPI_INT, nullptr, 0, MPI_INT, target, MPI_COMM_WORLD);
	}
	
	template<class T, class LocalData>
	void scatter(std::vector<T> & buffer, LocalData & lData)
	{
		MPI_Scatter(&buffer[0], 1, MPI_INT, &lData, 1, MPI_INT, 0, MPI_COMM_WORLD);
	}
	
	template<class S>
	void gather(S const& sendData, std::size_t root)
	{
		ENFORCE(MPI_Gather(const_cast<S*>(&sendData), 1, MPI_INT, nullptr, 0, MPI_INT, root, MPI_COMM_WORLD));
	}
	
	template<class S, class T>
	void gather(S const& sendData, std::vector<T> & buffer)
	{
		ENFORCE(MPI_Gather(const_cast<S*>(&sendData), 1, MPI_INT, &buffer[0], 1, MPI_INT, this->getProcessNumber(), MPI_COMM_WORLD));
	}*/
};

int main(int argc, char * argv[])
{	
	MPI_Encapsulation mpi(argc, &argv);
	
	std::vector<int> vec(mpi.getProcessCount(), 0);
	
	if ( mpi.isParentProcess() )
	{
		std::cout << "Starts with " << mpi.getProcessCount() << " Processes. (" << mpi.getProcessNumber() << ")" << std::endl;
		
		//vec.resize(mpi.getProcessCount(), 0);
		
		for(std::size_t i = 0; i < vec.size(); ++i)
			std::cout << " vec[" << i << "] = " << vec[i];
			
		std::cout << std::endl;
	}
	
	int & local = vec[mpi.getProcessNumber()];
	
	ENFORCE(MPI_Scatter(&vec[0], 1, MPI_INT, &local, 1, MPI_INT, 0, MPI_COMM_WORLD));
	
	local = mpi.getProcessNumber() + 1;
	
	ENFORCE(MPI_Gather(&local, 1, MPI_INT, &vec[0], 1, MPI_INT, 0, MPI_COMM_WORLD));
	
	
	if ( mpi.isParentProcess() )
	{
		bool success = true;
		for(std::size_t i = 0; i < vec.size(); ++i)
		{
			std::cout << " " << i << ": " << vec[i];
			if ( vec[i] != i+1 )
				success = false;
		}
		std::cout << "\nFinished with " << (success ? "success" : "fails") << std::endl;
		
		std::cout << std::endl;
	}

	return 0;
}
